/*******************************************************************************
 * Copyright (c) 2011 - 2012 TXT e-solutions SpA
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * This work was performed within the IoT_at_Work Project
 * and partially funded by the European Commission's
 * 7th Framework Programme under the research area ICT-2009.1.3
 * Internet of Things and enterprise environments.
 *
 *
 * Authors:
 *     Cristoforo Seccia (TXT e-solutions SpA)
 *
 * Contributors:
 *      Domenico Rotondi (TXT e-solutions SpA)
 *******************************************************************************/
package it.txt.access.capability.commons.signer;

import it.txt.access.capability.commons.osgi.CapabilityCommonsActivator;
import it.txt.access.capability.commons.signer.model.X509CertificateKeyValues;
import it.txt.access.capability.commons.signer.model.X509CertificateSubjectInfo;
import it.txt.access.capability.commons.signer.model.X509DocumentSignerInfo;
import it.txt.access.capability.commons.signer.utils.KeystoreHelper;

import java.io.FileInputStream;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyException;
import java.security.KeyPair;
import java.security.KeyStore;
import java.security.Provider;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.text.MessageFormat;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.xml.crypto.AlgorithmMethod;
import javax.xml.crypto.KeySelector;
import javax.xml.crypto.KeySelectorException;
import javax.xml.crypto.KeySelectorResult;
import javax.xml.crypto.MarshalException;
import javax.xml.crypto.XMLCryptoContext;
import javax.xml.crypto.XMLStructure;
import javax.xml.crypto.dsig.CanonicalizationMethod;
import javax.xml.crypto.dsig.DigestMethod;
import javax.xml.crypto.dsig.Reference;
import javax.xml.crypto.dsig.SignatureMethod;
import javax.xml.crypto.dsig.SignedInfo;
import javax.xml.crypto.dsig.Transform;
import javax.xml.crypto.dsig.XMLSignature;
import javax.xml.crypto.dsig.XMLSignatureException;
import javax.xml.crypto.dsig.XMLSignatureFactory;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.KeyValue;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;

import org.apache.xml.security.utils.ClassLoaderUtils;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

/**
 *
 * @author Cristoforo Seccia (TXT e-solutions SpA)
 * @author Will Provost
 * @author Salvatore Piccione (TXT e-solutions SpA - salvatore.piccione AT txtgroup.com)
 */
public class X509DocumentSigner {

    private static final Logger LOGGER = Logger.getLogger(X509DocumentSigner.class.getName());
    private static final String ERROR_RECOPVERING_KEYSTYORE = "Error while recovering keystore file";
    private static final String ERROR_CERT_ALIAS_NOT_FOUND = "Cannot find a certificate with the given signerID (%1$s)";
    private static final String ERROR_RECOVERING_CERT = "Error while searching for certificate into keystore";
    private static final String ERROR_RECOVERING_CERT_ALIAS = "Cannot recover the alias of the certificate";
    public static final String DEFAULT_SIGNATURE_PROVIDER_CLASS = "org.apache.jcp.xml.dsig.internal.dom.XMLDSigRI";
    
    public static final String SIGNATURE_PROVIDER_PROPERTY = "jsr105Provider";
    public static final String SIGNATURE_MECHANISM_TYPE = "DOM";
//    private static final ResourceBundle MESSAGES = ResourceBundle.getBundle(
//            "message-bundles/" + X509DocumentSigner.class.getSimpleName(), Locale.ROOT);
    private static XMLSignatureFactory factory;

    public static X509CertificateKeyValues getCertificateKeyValues(final X509DocumentSignerInfo model)
            throws GeneralSecurityException {
        //Initilixe the keystore
        KeyStore keyStore = null;
        try {
            //Recover the Keystore.
            keyStore = KeystoreHelper.getKeyStore(model.getKeystorePath(), model.getPrivateKeyPwd());
        } catch (GeneralSecurityException ex) {
            throw new GeneralSecurityException(ERROR_RECOPVERING_KEYSTYORE, ex);
        }
        //Initilixe the certificate
        java.security.cert.X509Certificate x509Certificate = null;
        try {
            //Recover the certificate from the Keystore
            x509Certificate = KeystoreHelper.getX509CertificateByIssuer(keyStore, model.getSignerID());
        } catch (GeneralSecurityException ex) {
            throw new GeneralSecurityException(ERROR_RECOVERING_CERT, ex);
        }
        //Check if we have found a certificate with the provided issuer.
        if (x509Certificate == null) {
            throw new GeneralSecurityException(String.format(ERROR_CERT_ALIAS_NOT_FOUND,model.getSignerID()));
        }
        try {
            //Recover the alias of the certificate
            String certificateAlias = KeystoreHelper.getAlias(keyStore, x509Certificate);
            //Recover the private key of the.
            X509CertificateKeyValues certificateKeyValues =
                    getCertificateKeyValues(
                    model.getKeystorePath(),
                    model.getPrivateKeyPwd(),
                    certificateAlias,
                    model.getPrivateKeyPwd());
            return certificateKeyValues;
        } catch (GeneralSecurityException ex) {
            throw new GeneralSecurityException(ERROR_RECOVERING_CERT_ALIAS, ex);
        }
    }

    public static X509CertificateKeyValues getCertificateKeyValues(
            final String keystorePath, final char[] keystorePsw,
            final String alias, final char[] certificatePsw) throws GeneralSecurityException {

        try {
            X509CertificateKeyValues keyValues = null;
            KeyStore keyStore = KeystoreHelper.getKeyStore(new FileInputStream(keystorePath), keystorePsw);
            KeyStore.PasswordProtection pp = new KeyStore.PasswordProtection(certificatePsw);
            KeyStore.PrivateKeyEntry entry = (KeyStore.PrivateKeyEntry) keyStore.getEntry(alias, pp);
            KeyPair keyPair = new KeyPair(entry.getCertificate().getPublicKey(), entry.getPrivateKey());

            KeyInfoFactory kFactory = getXMLSignatureFactory().getKeyInfoFactory();
            List<Certificate> certificates = Collections.singletonList(entry.getCertificate());
            X509Data x509Data = kFactory.newX509Data(certificates);
            List<X509Data> x509Datas = Collections.singletonList(x509Data);
            KeyInfo keyInfo = kFactory.newKeyInfo(x509Datas);

            keyValues = new X509CertificateKeyValues();
            keyValues.setKeyPair(keyPair);
            keyValues.setKeyInfo(keyInfo);
            keyValues.setAlias(alias);
            return keyValues;
        } catch (Exception ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        }
    }

    public static void signXMLElement(Element element, final X509CertificateKeyValues keyValues,
            final String afterNode) throws GeneralSecurityException {
        //check if we are in OSGi
        ClassLoader defaultClassLoader = null;
        boolean inOsgiContext = false;
        try {
            if (CapabilityCommonsActivator.getContextClassLoader() != null) {
                defaultClassLoader = Thread.currentThread().getContextClassLoader();
                inOsgiContext = true;
                if (LOGGER.isLoggable(Level.FINEST))
                    LOGGER.finest("Discarding default class loading because we are in an OSGi environment.");            
                Thread.currentThread().setContextClassLoader(CapabilityCommonsActivator.getContextClassLoader());
                if (LOGGER.isLoggable(Level.FINER))
                    LOGGER.finer("Using CapabilityCommonsActivator Class loader...");
            }
        } catch (NoClassDefFoundError e) {
            //maybe we are not in an OSGi context...
            if (!(e.getMessage().equals("it/txt/access/capability/commons/osgi/CapabilityCommonsActivator") ||
                    (e.getCause() instanceof ClassNotFoundException && "org.osgi.framework.BundleActivator".equals(e.getCause().getMessage()))))
                throw e;
        }
        
        try {

            String referenceUri = "";

            DigestMethod dm = getXMLSignatureFactory().newDigestMethod(DigestMethod.SHA1, null);

            Transform envelopingTransform;
            envelopingTransform = getXMLSignatureFactory().newTransform(Transform.ENVELOPED, (TransformParameterSpec) null);
//            if (envelopingTransform == null) {
//                LOGGER.log(Level.WARNING, "An error occurred while initialising the implementation of " + Transform.ENVELOPED);
//                envelopingTransform = loadTransformationImplementation();
//            } else
//                LOGGER.log(Level.INFO, "Despite of the runtime exception, this is the transformation: " + envelopingTransform);
            
            Reference ref = getXMLSignatureFactory().newReference(referenceUri, dm, 
                    Arrays.asList(envelopingTransform), null, null);

            CanonicalizationMethod method =
                    getXMLSignatureFactory().newCanonicalizationMethod(
                    //Canonicalization method: C14 INCLUSIVE WITH COMMENTS
                    CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS,
                    (C14NMethodParameterSpec) null);

            //TODO dynamic selection of the signature method
            String pkAlgorithm = keyValues.getKeyPair().getPrivate().getAlgorithm();
            SignatureMethod signatureMethod;
            if (pkAlgorithm.equals("DSA")) {
                signatureMethod = getXMLSignatureFactory().newSignatureMethod(
                        SignatureMethod.DSA_SHA1, null);
            }
            else if (pkAlgorithm.equals("RSA")) {
                signatureMethod = getXMLSignatureFactory().newSignatureMethod(
                        SignatureMethod.RSA_SHA1, null);
            } else
                throw new GeneralSecurityException("No signature method available for a " +
            		"private key generated with the algorithm " + pkAlgorithm);
            
            SignedInfo signedInfo = getXMLSignatureFactory().newSignedInfo(method, signatureMethod,
                    Collections.singletonList(ref));

            XMLSignature signature = getXMLSignatureFactory().newXMLSignature(signedInfo, keyValues.getKeyInfo());

            Node nextSibling = getNextSiblingNode(element, afterNode);

            DOMSignContext signContext = new DOMSignContext(keyValues.getKeyPair().getPrivate(),
                    element, nextSibling);
            
            //DEBUG - enable the collection of the digest process result for debugging purposes
//            signContext.setProperty("javax.xml.crypto.dsig.cacheReference", Boolean.TRUE);
            
            signature.sign(signContext);
            
            //Debug the signature process result
//            printReference(signature.getSignedInfo().getReferences().iterator());

        } catch (Exception ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        } finally {
            if (inOsgiContext && defaultClassLoader != null) {
                //reset the class loader
                Thread.currentThread().setContextClassLoader(defaultClassLoader);
            }
        }

    }

    private static DOMValidateContext getElementSignatureContext(Element element) throws GeneralSecurityException{
        NodeList nl = element.getElementsByTagNameNS(XMLSignature.XMLNS, "Signature");
        if (LOGGER.isLoggable(Level.FINEST))
            LOGGER.finest(MessageFormat.format("Total \"Signature\" elements found:{0}", nl.getLength()));
        if (nl.getLength() == 0) {
            throw new GeneralSecurityException("Cannot find Signature element");
        }
        //DEBUG - Print the content of the signature element
//        System.out.println("START Signature elements");
//        for (int i = 0; i < nl.getLength(); i++ ) {
//            System.out.println("Pretty printed signature " + i + ":: " + XMLPrinter.printDocumentNode(nl.item(i), false));
//        }
//        System.out.println("STOP Signature elements");
        return new DOMValidateContext(new KeyValueKeySelector(), nl.item(0));
    }

    @SuppressWarnings("unchecked")
    public static boolean verifyXMLElementSign(Element element) throws GeneralSecurityException {
        try {
            DOMValidateContext context = getElementSignatureContext(element);
            
            //DEBUG - enable the collection of the digest process result for debugging purposes
//            context.setProperty("javax.xml.crypto.dsig.cacheReference", Boolean.TRUE);
            
            // unmarshal the XMLSignature
            XMLSignature signature = getXMLSignatureFactory().unmarshalXMLSignature(context);
            
            // Validate the XMLSignature (generated above)
            boolean coreValidity = signature.validate(context);
            
            //Debug the signature process result
//            printReference(signature.getSignedInfo().getReferences().iterator());
            
            if (coreValidity) {
                if (LOGGER.isLoggable(Level.FINE))
                    LOGGER.fine("Signature passed core validation!");
                return true;
            } else {
                boolean sv = 
                        signature.getSignatureValue().validate(context);
                LOGGER.log(sv ? Level.FINER : Level.SEVERE, "The signature value is" + (sv ? " " : " NOT ") + "valid");
                Iterator<Reference> i = signature.getSignedInfo().getReferences().iterator();
                Reference ref;
                for (int j=0; i.hasNext(); j++) {
                    ref = i.next();
                    boolean refValid = ref.validate(context);
                    if (refValid) {
                        if (LOGGER.isLoggable(Level.FINER))
                            LOGGER.finer("the reference[" + j + "] is valid.");
                    }
                    else
                        LOGGER.log(Level.SEVERE, "the reference[" + j + "] is NOT valid.");
                }
                return false;
            }
        } catch (XMLSignatureException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        } catch (MarshalException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        } catch (NullPointerException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        }
    }
    
//    private static void printReference (Iterator<Reference> referenceIterator) throws IOException {
//        Reference currentRef;
//        InputStream is;
//        InputStreamReader isr;
//        StringBuilder buffer;
//        int currentCharAsInt;
//        while (referenceIterator.hasNext()) {
//            currentRef = referenceIterator.next();
//            is = currentRef.getDigestInputStream();
//            if (is == null) {
//                LOGGER.log(Level.WARNING, "No input stream for the current reference");
//                continue;
//            }
//            isr = new InputStreamReader(is, Charset.forName("UTF-8"));
//            buffer = new StringBuilder();
//            while ((currentCharAsInt = isr.read()) != -1) {
//                buffer.append((char)currentCharAsInt);
//            }
//            LOGGER.log(Level.INFO, "Value analysed::\n" + buffer.toString());
//        }
//    }

    public static boolean verifyXMLElementSign(Element element,
            X509CertificateSubjectInfo subjectInfo) throws GeneralSecurityException {
        try {
            DOMValidateContext context = getElementSignatureContext(element);
            XMLSignature signature = getXMLSignatureFactory().unmarshalXMLSignature(context);
            if (signature.validate(context)) {
                if (LOGGER.isLoggable(Level.FINER)) {
                    LOGGER.finer("Signature passed core validation. " +
                    		"Recovering X509 certificate public values...");
                }
                if (subjectInfo == null) {
                    subjectInfo = new X509CertificateSubjectInfo();
                }
                X509CertificateSubjectInfo info = getX509CertificateSubjectInfo(signature);
                subjectInfo.setX509IssuerName(info.getX509IssuerName());
                subjectInfo.setX509SerialNumber(info.getX509SerialNumber());
                subjectInfo.setX509SubjectName(info.getX509SubjectName());
                if (LOGGER.isLoggable(Level.FINER)) {
                    LOGGER.finer(MessageFormat.format("Found Ceritificate Subject DN: {0}",
                            subjectInfo.getX509SubjectName()));
                    LOGGER.finer(MessageFormat.format("Found Ceritificate Issuer DN: {0}",
                            subjectInfo.getX509IssuerName()));
                    LOGGER.finer(MessageFormat.format("Found Ceritificate Serial Number: {0}",
                            subjectInfo.getX509SerialNumber().toString()));
                }
                return true;
            } else {
                LOGGER.log(Level.SEVERE, "Signature failed core validation");
                return false;
            }
        } catch (XMLSignatureException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        } catch (MarshalException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        } catch (NullPointerException ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        }
    }

    private static X509CertificateSubjectInfo getX509CertificateSubjectInfo(final XMLSignature signature) {
        X509CertificateSubjectInfo x509SubjectInfo = new X509CertificateSubjectInfo();
        for (Object keyInfoItem : signature.getKeyInfo().getContent()) {
            if (keyInfoItem instanceof X509Data) {
                for (Object X509Item : ((X509Data) keyInfoItem).getContent()) {
                    if (X509Item instanceof X509Certificate) {
                        X509Certificate xCert = (X509Certificate) X509Item;
                        x509SubjectInfo.setX509SubjectName(xCert.getSubjectX500Principal().toString());
                        x509SubjectInfo.setX509IssuerName(xCert.getIssuerDN().toString());
                        x509SubjectInfo.setX509SerialNumber(xCert.getSerialNumber());
                        break;
                    }
                }
                break;
            }
        }
        return x509SubjectInfo;
    }

    private static XMLSignatureFactory newXMLSignatureFactory() throws GeneralSecurityException {
        try {
            if (LOGGER.isLoggable(Level.FINER)) {
                LOGGER.finer("XMLSignatureFactory - Using provider Property: " + SIGNATURE_PROVIDER_PROPERTY);
                LOGGER.finer("XMLSignatureFactory - Using Signature Mechanism Type: " + SIGNATURE_MECHANISM_TYPE);
            }
            String providerName = System.getProperty(SIGNATURE_PROVIDER_PROPERTY, DEFAULT_SIGNATURE_PROVIDER_CLASS);
            if (LOGGER.isLoggable(Level.FINER))
                LOGGER.finer("XMLSignatureFactory - Using Signature Provider Class: " + providerName);
            //In order to avoid a Caused by: java.lang.ClassNotFoundException: org.apache.jcp.xml.dsig.internal.dom.XMLDSigRI
            
            Provider provider = (Provider) ClassLoaderUtils.loadClass(DEFAULT_SIGNATURE_PROVIDER_CLASS,
                    X509DocumentSigner.class).newInstance();
//            Provider provider = (Provider) Class.forName(providerName).newInstance();
            return XMLSignatureFactory.getInstance(SIGNATURE_MECHANISM_TYPE, provider);
        } catch (Exception ex) {
            throw new GeneralSecurityException(ex.getMessage(), ex);
        }
    }

    private static XMLSignatureFactory getXMLSignatureFactory() throws GeneralSecurityException {
        if (factory == null) {
            if (LOGGER.isLoggable(Level.FINER))
                LOGGER.finer("Creating new XMLSignatureFactory");
            factory = newXMLSignatureFactory();
        }
        return factory;
    }

    /**
     * Adds an enveloped signature to the given element. The afterNode cannot be the last in the schema.
     */
    /**
     * Util method used to retrieve a node into an xml.
     * @return org.w3c.Node.
     */
    private static Node getNextSiblingNode(Element target, String nodeName) {

        boolean found = false;
        Node nextSibling = null;
        NodeList children = target.getChildNodes();
        for (int c = 0; c < children.getLength(); ++c) {
            Node child = children.item(c);
            if (found) {
                nextSibling = child;
                break;
            }
            if (child.getNodeType() == Node.ELEMENT_NODE
                    && child.getLocalName().equals(nodeName)) {
                found = true;
            }
        }
        return nextSibling;
    }

    /**
     * KeySelector that can handle KeyValue and X509Data info.
     */
    private static class KeyValueKeySelector extends KeySelector {

        public KeySelectorResult select(KeyInfo keyInfo,
                KeySelector.Purpose purpose,
                AlgorithmMethod method,
                XMLCryptoContext context) throws KeySelectorException {

            if (keyInfo == null) {
                throw new KeySelectorException("Null KeyInfo object!");
            }

            SignatureMethod sm = (SignatureMethod) method;

            @SuppressWarnings("unchecked")
            List<XMLStructure> list = keyInfo.getContent();
            XMLStructure xmlStructure;
            for (int i = 0; i < list.size(); i++) {
                xmlStructure = list.get(i);
                PublicKey pk = null;

                try {
                    if (xmlStructure instanceof KeyValue) {
                        pk = ((KeyValue) xmlStructure).getPublicKey();
                    } else if (xmlStructure instanceof X509Data) {
                        for (Object data :
                                ((X509Data) xmlStructure).getContent()) {
                            if (data instanceof X509Certificate) {
                                pk = ((X509Certificate) data).getPublicKey();
                            }
                        }
                    }
                } catch (KeyException ke) {
                    throw new KeySelectorException(ke);
                }

                if (algEquals(sm.getAlgorithm(), pk.getAlgorithm())) {
                    return new SimpleKeySelectorResult(pk);
                }
            }
            throw new KeySelectorException("No KeyValue element found!");
        }
    }

    /**
     * Test that a formal URI expresses the same algorithm as a conventional
     * short name such as "DSA" or "RSA".
     */
    private static boolean algEquals(String algURI, String algName) {
        return (algName.equalsIgnoreCase("DSA")
                && algURI.equalsIgnoreCase(SignatureMethod.DSA_SHA1))
                || (algName.equalsIgnoreCase("RSA")
                && algURI.equalsIgnoreCase(SignatureMethod.RSA_SHA1));
    }

    /**
     * Data structure returned by the key selector to the validation context.
     */
    private static class SimpleKeySelectorResult implements KeySelectorResult {

        private PublicKey pk;

        SimpleKeySelectorResult(PublicKey pk) {
            this.pk = pk;
        }

        public Key getKey() {
            return pk;
        }
    }
    
    public static X509CertificateSubjectInfo getX509CertificateSubjectInfo (Element xmlElement) throws GeneralSecurityException, MarshalException {
        DOMValidateContext context = getElementSignatureContext(xmlElement);
        XMLSignature signature = getXMLSignatureFactory().unmarshalXMLSignature(context);
        X509CertificateSubjectInfo info = getX509CertificateSubjectInfo(signature);
        X509CertificateSubjectInfo subjectInfo = new X509CertificateSubjectInfo();
        subjectInfo.setX509IssuerName(info.getX509IssuerName());
        subjectInfo.setX509SerialNumber(info.getX509SerialNumber());
        subjectInfo.setX509SubjectName(info.getX509SubjectName());
        return subjectInfo;
    }
    
//    private static Transform loadTransformationImplementation () throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, GeneralSecurityException {
//        ClassLoader defaultClassLoader = Thread.currentThread().getContextClassLoader();
//        if (LOGGER.isLoggable(Level.FINE))
//            LOGGER.log(Level.FINE, "Default initialization failed. Bundle initialization...");            
//        
//        if (CapabilityCommonsActivator.getContextClassLoader() != null) {
//            if (LOGGER.isLoggable(Level.FINE))
//                LOGGER.log(Level.FINE, "Using CapabilityCommonsActivator Class loader...");
//            Thread.currentThread().setContextClassLoader(CapabilityCommonsActivator.getContextClassLoader());
//        } else {
//            if (LOGGER.isLoggable(Level.FINE))
//                LOGGER.log(Level.FINE, "Using XMLPrinter Class loader...");
//            Thread.currentThread().setContextClassLoader(X509DocumentSigner.class.getClassLoader());
//        }
//        Transform envelopingTransform = getXMLSignatureFactory().newTransform(Transform.ENVELOPED, (TransformParameterSpec) null);
//        //reset the class loader
//        Thread.currentThread().setContextClassLoader(defaultClassLoader);
//        
//        return envelopingTransform;
//    }
}
